{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "模型的部署参考： [learn-llm-deploy-easily](https://gitee.com/coderwillyan/learn-llm-deploy-easily) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里主要介绍如何调用已部署的模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# embedding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用embedding API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以zhipu为例"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.embeddings import ZhipuAIEmbeddings\n",
    "\n",
    "my_emb = ZhipuAIEmbeddings(\n",
    "    model=\"embedding-2\",\n",
    "    api_key=os.environ[\"ZHIPUAI_API_KEY\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用本地部署的embedding模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hugging Face Transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "​​特点​​：底层灵活，支持BERT、RoBERTa等Transformer架构的Embedding模型，需自定义向量提取逻辑"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from transformers import AutoModel, AutoTokenizer\n",
    "# model = AutoModel.from_pretrained(\"BAAI/bge-large-zh\")\n",
    "# embeddings = model(inputs).last_hidden_state.mean(dim=1)  # 提取句向量"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sentence Transformers部署"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "​​特点​​：基于Transformers的封装，内置池化层，一键生成句子级嵌入，支持语义相似度计算"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "项目地址： https://github.com/UKPLab/sentence-transformers\n",
    "\n",
    "参考文档： https://sbert.net.cn/index.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'sentence_transformers'",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mModuleNotFoundError\u001b[39m                       Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 4\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;66;03m# Requires transformers>=4.51.0\u001b[39;00m\n\u001b[32m      2\u001b[39m \u001b[38;5;66;03m# Requires sentence-transformers>=2.7.0\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m4\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01msentence_transformers\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m SentenceTransformer\n\u001b[32m      6\u001b[39m \u001b[38;5;66;03m# Load the model\u001b[39;00m\n\u001b[32m      7\u001b[39m model = SentenceTransformer(\u001b[33m\"\u001b[39m\u001b[33m/opt/workspace/models/Qwen/Qwen3-Embedding-0.6B\u001b[39m\u001b[33m\"\u001b[39m,  device=\u001b[33m\"\u001b[39m\u001b[33mcuda:3\u001b[39m\u001b[33m\"\u001b[39m)\n",
      "\u001b[31mModuleNotFoundError\u001b[39m: No module named 'sentence_transformers'"
     ]
    }
   ],
   "source": [
    "# Requires transformers>=4.51.0\n",
    "# Requires sentence-transformers>=2.7.0\n",
    "\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "# Load the model\n",
    "model = SentenceTransformer(\"/opt/workspace/models/Qwen/Qwen3-Embedding-0.6B\",  device=\"cuda:3\")\n",
    "\n",
    "# We recommend enabling flash_attention_2 for better acceleration and memory saving,\n",
    "# together with setting `padding_side` to \"left\":\n",
    "# model = SentenceTransformer(\n",
    "#     \"Qwen/Qwen3-Embedding-0.6B\",\n",
    "#     model_kwargs={\"attn_implementation\": \"flash_attention_2\", \"device_map\": \"auto\"},\n",
    "#     tokenizer_kwargs={\"padding_side\": \"left\"},\n",
    "# )\n",
    "\n",
    "# The queries and documents to embed\n",
    "queries = [\n",
    "    \"What is the capital of China?\",\n",
    "    \"Explain gravity\",\n",
    "]\n",
    "documents = [\n",
    "    \"The capital of China is Beijing.\",\n",
    "    \"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.\",\n",
    "]\n",
    "\n",
    "# Encode the queries and documents. Note that queries benefit from using a prompt\n",
    "# Here we use the prompt called \"query\" stored under `model.prompts`, but you can\n",
    "# also pass your own prompt via the `prompt` argument\n",
    "query_embeddings = model.encode(queries, prompt_name=\"query\")\n",
    "document_embeddings = model.encode(documents)\n",
    "\n",
    "# Compute the (cosine) similarity between the query and document embeddings\n",
    "similarity = model.similarity(query_embeddings, document_embeddings)\n",
    "print(similarity)\n",
    "# tensor([[0.7646, 0.1414],\n",
    "#         [0.1355, 0.6000]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "# 指定 GPU 设备（例如 GPU 0）\n",
    "# device = \"cuda:3\" if torch.cuda.is_available() else \"cpu\"\n",
    "# Load the model\n",
    "model = SentenceTransformer(\"/opt/workspace/models/Qwen/Qwen3-Embedding-0.6B\", device=\"cuda:3\")\n",
    "result = model.encode(\"这是一个测试\")\n",
    "result.tolist()[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### HuggingFaceEmbeddings部署"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings\n",
    "my_emb = HuggingFaceEmbeddings(model_name='/opt/workspace/models/Qwen/Qwen3-Embedding-0.6B', model_kwargs={\"device\": \"cuda:3\"})\n",
    "query = \"如何调用HuggingFaceEmbeddings？\"\n",
    "query_vector = my_emb.embed_query(query) \n",
    "print(\"查询向量:\", query_vector[:5], \"...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ollama部署"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from langchain_community.embeddings import OllamaEmbeddings\n",
    "from langchain_ollama.embeddings import OllamaEmbeddings\n",
    "my_emb = OllamaEmbeddings(base_url='http://localhost:11434',model=\"dengcao/Qwen3-Embedding-0.6B:F16\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "查询向量: [-0.01782775, -0.010210706, -0.031165373, 0.029789878, -0.024330704] ...\n"
     ]
    }
   ],
   "source": [
    "from langchain_ollama.embeddings import OllamaEmbeddings\n",
    "\n",
    "# 初始化嵌入模型\n",
    "my_emb = OllamaEmbeddings(\n",
    "    base_url='http://localhost:11434',\n",
    "    model=\"bge-m3:latest\"\n",
    ")\n",
    "\n",
    "# 示例文本\n",
    "query = \"如何调用OllamaEmbeddings？\"\n",
    "\n",
    "# 生成嵌入向量\n",
    "query_vector = my_emb.embed_query(query)\n",
    "\n",
    "# 输出结果\n",
    "print(\"查询向量:\", query_vector[:5], \"...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### vllm部署"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 08-19 10:30:38 [utils.py:326] non-default args: {'model': '/workspace/models/Qwen/Qwen3-Embedding-0___6B', 'task': 'embed', 'dtype': 'float16', 'disable_log_stats': True, 'enforce_eager': True}\n",
      "INFO 08-19 10:30:38 [__init__.py:711] Resolved architecture: Qwen3ForCausalLM\n",
      "WARNING 08-19 10:30:38 [__init__.py:2819] Casting torch.bfloat16 to torch.float16.\n",
      "INFO 08-19 10:30:38 [__init__.py:1750] Using max model len 32768\n",
      "WARNING 08-19 10:30:38 [cuda.py:103] To see benefits of async output processing, enable CUDA graph. Since, enforce-eager is enabled, async output processor cannot be used\n",
      "INFO 08-19 10:30:38 [__init__.py:3565] Cudagraph is disabled under eager mode\n",
      "INFO 08-19 10:30:38 [llm_engine.py:222] Initializing a V0 LLM engine (v0.10.1) with config: model='/workspace/models/Qwen/Qwen3-Embedding-0___6B', speculative_config=None, tokenizer='/workspace/models/Qwen/Qwen3-Embedding-0___6B', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config={}, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=32768, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=True, kv_cache_dtype=auto, device_config=cuda, decoding_config=DecodingConfig(backend='auto', disable_fallback=False, disable_any_whitespace=False, disable_additional_properties=False, reasoning_backend=''), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None), seed=None, served_model_name=/workspace/models/Qwen/Qwen3-Embedding-0___6B, enable_prefix_caching=None, chunked_prefill_enabled=False, use_async_output_proc=False, pooler_config=PoolerConfig(pooling_type='LAST', normalize=True, dimensions=None, activation=None, softmax=None, step_tag_id=None, returned_token_ids=None, enable_chunked_processing=None, max_embed_len=None), compilation_config={\"level\":0,\"debug_dump_path\":\"\",\"cache_dir\":\"\",\"backend\":\"\",\"custom_ops\":[],\"splitting_ops\":null,\"use_inductor\":true,\"compile_sizes\":[],\"inductor_compile_config\":{\"enable_auto_functionalized_v2\":false},\"inductor_passes\":{},\"cudagraph_mode\":0,\"use_cudagraph\":true,\"cudagraph_num_of_warmups\":0,\"cudagraph_capture_sizes\":[],\"cudagraph_copy_inputs\":false,\"full_cuda_graph\":false,\"pass_config\":{},\"max_capture_size\":0,\"local_cache_dir\":null}, use_cached_outputs=False, \n",
      "INFO 08-19 10:30:39 [model_runner.py:1080] Starting to load model /workspace/models/Qwen/Qwen3-Embedding-0___6B...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  3.81it/s]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  3.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 08-19 10:30:40 [default_loader.py:262] Loading weights took 0.29 seconds\n",
      "INFO 08-19 10:30:40 [model_runner.py:1112] Model loading took 1.1098 GiB and 0.337335 seconds\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 08-19 10:30:40 [llm.py:298] Supported_tasks: ['encode', 'embed']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Adding requests: 100%|██████████| 4/4 [00:00<00:00, 2938.22it/s]\n",
      "Processed prompts: 100%|██████████| 4/4 [00:00<00:00, 80.88it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.7643852829933167, 0.1411387324333191], [0.13565093278884888, 0.5996948480606079]]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Requires vllm>=0.8.5\n",
    "import torch\n",
    "import vllm\n",
    "from vllm import LLM\n",
    "\n",
    "def get_detailed_instruct(task_description: str, query: str) -> str:\n",
    "    return f'Instruct: {task_description}\\nQuery:{query}'\n",
    "\n",
    "# Each query must come with a one-sentence instruction that describes the task\n",
    "task = 'Given a web search query, retrieve relevant passages that answer the query'\n",
    "\n",
    "queries = [\n",
    "    get_detailed_instruct(task, 'What is the capital of China?'),\n",
    "    get_detailed_instruct(task, 'Explain gravity')\n",
    "]\n",
    "# No need to add instruction for retrieval documents\n",
    "documents = [\n",
    "    \"The capital of China is Beijing.\",\n",
    "    \"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.\"\n",
    "]\n",
    "input_texts = queries + documents\n",
    "\n",
    "# model = LLM(model=\"/opt/workspace/models/Qwen/Qwen3-Embedding-0.6B\", task=\"embed\")\n",
    "\n",
    "model = LLM(\n",
    "    model=\"/workspace/models/Qwen/Qwen3-Embedding-0___6B\",\n",
    "    task=\"embed\",\n",
    "    # trust_remote_code=True,  \n",
    "    enforce_eager=True,      # 避免可能的图优化问题\n",
    "    dtype=\"float16\" \n",
    ")\n",
    "\n",
    "outputs = model.embed(input_texts)\n",
    "embeddings = torch.tensor([o.outputs.embedding for o in outputs])\n",
    "scores = (embeddings[:2] @ embeddings[2:].T)\n",
    "print(scores.tolist())\n",
    "# [[0.7620252966880798, 0.14078938961029053], [0.1358368694782257, 0.6013815999031067]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Docker方式部署"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### xinference部署"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.embeddings import XinferenceEmbeddings\n",
    "\n",
    "# 替换为你的Xinference服务器URL和模型UID\n",
    "xinference = XinferenceEmbeddings(\n",
    "    server_url=\"http://localhost:9997\",  # 注意：原代码中的\"loaclhost\"拼写错误，应为\"localhost\"\n",
    "    model_uid=\"your_model_uid\"  # 替换为实际的模型UID\n",
    ")\n",
    "\n",
    "# 输入文本\n",
    "texts = [\"你好，世界。\", \"LangChain 是一个强大的工具。\"]\n",
    "\n",
    "# 生成嵌入向量\n",
    "vectors = xinference.embed_documents(texts)\n",
    "\n",
    "# 打印结果\n",
    "for idx, vector in enumerate(vectors):\n",
    "    print(f\"文本 {idx + 1}: {texts[idx]}\")\n",
    "    print(f\"嵌入向量: {vector[:5]}... (维度: {len(vector)})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3 (env_vllm)",
   "language": "python",
   "name": "env_vllm"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
